import pandas as pd
from .utils.Distribution import *
from .utils.Compute import *
from .utils.Save import *
from .utils.Draw import *

# 保留 generate_demographics_df 及其依赖函数

def generate_demographics_df(
    n_sample,
    female_ratio,
    low_age,
    up_age,
    male_distri,
    female_distri,
    p_below_65_male=0.5,
    p_below_65_female=0.5,
    c0_w_m=0, c1_w_m=0, c0_w_f=0, c1_w_f=0,
    c0_h_m=0, c1_h_m=0, c2_h_m=0, c0_h_f=0, c1_h_f=0, c2_h_f=0,
    c1_param=0, weight_ex=0, height_ex=0,
    alpha_male=1.5, beta_male=None, alpha_female=1.5, beta_female=None,
    cv_male_weight=0, cv_female_weight=0, cv_male_height=0, cv_female_height=0,
    seed=42
):
    male_ratio = 1 - female_ratio
    df = generate_population(
        male_distribution=male_distri,
        female_distribution=female_distri,
        male_ratio=male_ratio,
        n_samples=n_sample,
        seed=seed,
        low_age=low_age,
        up_age=up_age,
        alpha_male=alpha_male,
        beta_male=beta_male,
        alpha_female=alpha_female,
        beta_female=beta_female,
        p_below_65_male=p_below_65_male,
        p_below_65_female=p_below_65_female
    )
    age_tuples, height_tuples, weight_tuples, bsa_tuples, id_tuples = process_population_data(
        df,
        c0_w_m, c1_w_m,
        c0_w_f, c1_w_f,
        c0_h_m, c1_h_m, c2_h_m,
        c0_h_f, c1_h_f, c2_h_f,
        c1_param, weight_ex, height_ex,
        cv_male_height, cv_male_weight,
        cv_female_height, cv_female_weight,
        seed
    )
    # 男性数据
    ages_male, ages_female = age_tuples
    heights_male, heights_female = height_tuples
    weights_male, weights_female = weight_tuples
    bsas_male, bsas_female = bsa_tuples
    ids_male, ids_female = id_tuples
    male_df = pd.DataFrame({
        'id': ids_male,
        'gender': 'male',
        'age': ages_male,
        'height': heights_male,
        'weight': weights_male,
        'bsa': bsas_male
    }) if len(ids_male) > 0 else pd.DataFrame()
    female_df = pd.DataFrame({
        'id': ids_female,
        'gender': 'female',
        'age': ages_female,
        'height': heights_female,
        'weight': weights_female,
        'bsa': bsas_female
    }) if len(ids_female) > 0 else pd.DataFrame()
    processed_df = pd.concat([male_df, female_df], ignore_index=True)
    return processed_df

# --- 程序入口 ---
if __name__ == "__main__":
    # 示例调用，用于测试
    n_sample = 1000
    female_ratio = 0.5
    low_age = 20
    up_age = 80
    male_distri = 'Uniform'
    female_distri = 'Uniform'
    p_below_65_male = 0.5
    p_below_65_female = 0.5
    c0_w_m = 0
    c1_w_m = 0
    c0_w_f = 0
    c1_w_f = 0
    c0_h_m = 0
    c1_h_m = 0
    c2_h_m = 0
    c0_h_f = 0
    c1_h_f = 0
    c2_h_f = 0
    c1_param = 0
    weight_ex = 0
    height_ex = 0
    alpha_male = 1.5
    beta_male = None
    alpha_female = 1.5
    beta_female = None
    cv_male_weight = 0
    cv_female_weight = 0
    cv_male_height = 0
    cv_female_height = 0
    seed = 42

    processed_df = generate_demographics_df(
        n_sample=n_sample,
        female_ratio=female_ratio,
        low_age=low_age,
        up_age=up_age,
        male_distri=male_distri,
        female_distri=female_distri,
        p_below_65_male=p_below_65_male,
        p_below_65_female=p_below_65_female,
        c0_w_m=c0_w_m, c1_w_m=c1_w_m, c0_w_f=c0_w_f, c1_w_f=c1_w_f,
        c0_h_m=c0_h_m, c1_h_m=c1_h_m, c2_h_m=c2_h_m, c0_h_f=c0_h_f, c1_h_f=c1_h_f, c2_h_f=c2_h_f,
        c1_param=c1_param, weight_ex=weight_ex, height_ex=height_ex,
        alpha_male=alpha_male, beta_male=beta_male, alpha_female=alpha_female, beta_female=beta_female,
        cv_male_weight=cv_male_weight, cv_female_weight=cv_female_weight, cv_male_height=cv_male_height, cv_female_height=cv_female_height,
        seed=seed
    )
    print("生成的人口数据:")
    print(processed_df.head())
    print(f"\n生成 {len(processed_df)} 条数据，其中男性 {len(processed_df[processed_df['gender'] == 'male'])}，女性 {len(processed_df[processed_df['gender'] == 'female'])}。\n")

    # 保存结果
    output_path = r"C:\Users\27135\Desktop\population\封装\src\module1"
    filename_prefix = f"{male_distri}_{female_distri}"
    print(f"正在保存结果到路径: {output_path}，文件名前缀: {filename_prefix}...")
    save_demographics_to_csv((age_tuples, height_tuples, weight_tuples, bsa_tuples, id_tuples), filename_prefix, output_path)
    print(f"数据成功保存!")

    # 绘制图形并保存
    print("正在生成人口分布图表...")
    fig = plot_population_histograms(
        (age_tuples, height_tuples, weight_tuples, bsa_tuples, id_tuples),
        n_sample,
        f"{male_distri}_{female_distri}",
        output_path=output_path,
        low_age=low_age,
        up_age=up_age
    )
    print("图表生成并保存完成!")
